package com.gallop.mq.core.dispatcher;

import com.gallop.mq.core.exception.MQException;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.gallop.mq.core.MQListenerInterceptor;
import com.gallop.mq.core.MQMessageDispatcher;
import com.gallop.mq.core.MQMessageListenerContainer;
import com.gallop.mq.core.MQReliableRegistry;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Slf4j
@Data
public class MQMessageDispatcherImpl implements MQMessageDispatcher {

    private final Map<String, MQReliableRegistry.ReliableDeliveryConfig> reliableConfigMap;
    private MQReliableRegistry reliableRegistry;

    public MQMessageDispatcherImpl(MQReliableRegistry reliableRegistry) {
        this.reliableRegistry = reliableRegistry;
        this.reliableConfigMap = reliableRegistry.register().stream()
                .collect(Collectors.toMap(MQReliableRegistry.ReliableDeliveryConfig::getTopic, c -> c));
    }

    private final Multimap<Class<?>, MQMessageListenerContainer> routes = Multimaps.synchronizedSetMultimap(HashMultimap.create());

    @Override
    public void addRoutes(Collection<MQMessageListenerContainer> containers) {
        containers.forEach(container -> this.routes.put(container.getMessageType(), container));
    }

    @Override
    public void dispatch(DispatchContext dispatchContext) {
        Object message = dispatchContext.getMessage();
        String group = dispatchContext.getGroup();
        String topic = dispatchContext.getTopic();
        if (message == null) {
            return;
        }
        Collection<MQMessageListenerContainer> routees = routes.keySet().stream()
                .filter(t -> t.isAssignableFrom(message.getClass()))
                .flatMap(t -> routes.get(t).stream()).collect(Collectors.toSet());
//        Collection<MQMessageListenerContainer> routees = routes.get(message.getClass());
        if (!StringUtils.hasLength(group)) {
            routees.stream().filter(container -> topic.equals(container.getTopic()))
                    .forEach(container -> doDispatch(dispatchContext, container));
        } else {
            routees.stream().filter(container -> topic.equals(container.getTopic()) && group.equals(container.getGroup()))
                    .findAny()
                    .ifPresent(container -> doDispatch(dispatchContext, container));
        }
    }

    private void doDispatch(DispatchContext dispatchContext, MQMessageListenerContainer container) {
        Object message = dispatchContext.getMessage();
        Class<?>[] paramTypes = container.getMethod().getParameterTypes();
        if (paramTypes.length <= 0) {
            throw new IllegalStateException("MQ container dispatching error, need at least 1 message param in method="
                    + container.getMethod() + ", container=" + container);
        }
        // 拼接调用参数
        Object[] args = buildInvocationArgs(dispatchContext, paramTypes);
        try {
            // 拦截器触发
            List<MQListenerInterceptor> interceptors = container.getInterceptor();
            interceptors.forEach(interceptor -> interceptor.beforeInvoke(message));
            // 判断调用方式 sub/rpc
            Object result = isRpcDispatch(dispatchContext) ? handleReply(dispatchContext, container, args) :
                    handleSubscription(dispatchContext, container, args);
            // 拦截器触发
            interceptors.forEach(interceptor -> interceptor.afterInvoke(message, result));
        } catch (
                Exception e) {
            log.error("MQ message dispatch failed.", e);
            List<MQListenerInterceptor> interceptors = container.getInterceptor();
            interceptors.forEach(interceptor -> interceptor.onError(message, e));
        }

    }

    private Object[] buildInvocationArgs(DispatchContext dispatchContext, Class<?>[] paramTypes) {
        return Arrays.stream(paramTypes).map(paramType -> {
            if (DispatchRpcCallable.class.isAssignableFrom(paramType)) {
                return isRpcDispatch(dispatchContext) ? DispatchRpcCallable.builder()
                        .replyTo(dispatchContext.getReplyTo())
                        .correlationId(dispatchContext.getCorrelationId())
                        .rpcResponse(dispatchContext.getRpcResponse())
                        .build() : DispatchRpcCallable.nop();
            }
            if (paramType.isAssignableFrom(dispatchContext.getMessage().getClass())) {
                return dispatchContext.getMessage();
            }
            return null;
        }).toArray();
    }

    private boolean isRpcDispatch(DispatchContext dispatchContext) {
        return StringUtils.hasLength(dispatchContext.getReplyTo()) && dispatchContext.getRpcResponse() != null;
    }

    private Object handleSubscription(DispatchContext dispatchContext, MQMessageListenerContainer container, Object[] args) throws Exception {
        if (!handleReliableAck(dispatchContext, true)) {
            throw new MQException(String.format("handleReliableAck before failed. topic=[%s],group=[%s],message=[%s]",
                    dispatchContext.getTopic(), dispatchContext.getGroup(), dispatchContext.getMessage().getClass()));
        }
        Object result = container.getMethod().invoke(container.getInstance(), args);
        if (handleReliableAck(dispatchContext, false)) {
            return result;
        }
        throw new MQException(String.format("handleReliableAck after failed. topic=[%s],group=[%s],message=[%s]",
                dispatchContext.getTopic(), dispatchContext.getGroup(), dispatchContext.getMessage().getClass()));
    }

    private Object handleReply(DispatchContext dispatchContext, MQMessageListenerContainer container, Object[] args) throws Exception {
        boolean callbackAsParam = Arrays.stream(args).anyMatch(arg -> DispatchRpcCallable.class.isAssignableFrom(arg.getClass()));
        Class<?> returnType = container.getMethod().getReturnType();

        boolean callbackAsReturn = !callbackAsParam && !returnType.equals(void.class);
        if (!callbackAsParam && !callbackAsReturn) {
            log.error("MQ container dispatching message in rpc mode but callback not found,fallback to subscription mode.method={}, container={}"
                    , container.getMethod(), container);
            return handleSubscription(dispatchContext, container, args);
        }
        Object result = container.getMethod().invoke(container.getInstance(), args);
        if (callbackAsReturn) {
            DispatchRpcCallable rpcCallable;
            if (DispatchRpcCallable.class.isAssignableFrom(returnType)) {
                // 使用监听方法返回值作为callback
                rpcCallable = (DispatchRpcCallable) result;
            } else {
                rpcCallable = DispatchRpcCallable.builder()
                        .replyTo(dispatchContext.getReplyTo())
                        .correlationId(dispatchContext.getCorrelationId())
                        .rpcResponse(dispatchContext.getRpcResponse())
                        .build();
            }
            rpcCallable.reply(result);
        }
        return result;
    }

    private boolean handleReliableAck(DispatchContext dispatchContext, boolean before) {
        String topic = dispatchContext.getTopic();
        String group = dispatchContext.getGroup();
        Object message = dispatchContext.getMessage();

        MQReliableRegistry.ReliableDeliveryConfig reliableCfg = this.reliableConfigMap.get(topic);
        if (reliableCfg != null && reliableCfg.getMessageType().isAssignableFrom(message.getClass()) &&
                reliableCfg.getListeners().contains(group)) {
            return before ? reliableCfg.getBeforeAck().apply(reliableCfg, topic, message, group) :
                    reliableCfg.getAfterAck().apply(reliableCfg, topic, message, group);
        }
        return true;
    }

    @Override
    public Collection<MQMessageListenerContainer> getAllContainers() {
        return routes.values();
    }
}
